import asyncio
import json
import logging
from typing import List, AsyncGenerator, Any, Dict

# 从您现有的 chat.py 模块中导入底层的异步流式生成函数
from chat import generate_chat_stream

logger = logging.getLogger(__name__)


# --- 1. (关键) "异步转同步"的桥接函数 (保持不变) ---
def stream_async_generator(async_gen: AsyncGenerator[Any, None]):
    """
    一个同步的生成器函数，用于包装和执行一个异步生成器。
    """
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    async_iterator = async_gen.__aiter__()
    try:
        while True:
            next_item = loop.run_until_complete(async_iterator.__anext__())
            yield next_item
    except StopAsyncIteration:
        pass


# --- 2. 异步的辅助函数 (关键改动：解析为嵌套结构) ---
async def call_llm_for_parsing(text: str) -> str:
    """
    一个辅助函数，用于调用LLM并获取完整的非流式响应。
    """
    full_response = ""

    prompt = f"""请分析以下报告大纲，并将其转换为结构化的JSON格式。
JSON应该是一个对象的数组。每个对象代表一个一级标题，并且必须包含两个键：
1. "heading": 一个字符串，包含一级标题的名称（例如："1. 行业概述"）。
2. "subheadings": 一个字符串数组，包含该一级标题下的所有二级子标题（例如：["1.1 领域定义与发展历程", "1.2 当前市场规模与主要参与者"]）。

如果某个一级标题下没有子标题，则"subheadings"数组应为空。

输出格式示例：
[
  {{
    "heading": "1. 行业概述",
    "subheadings": ["1.1 领域定义与发展历程", "1.2 当前市场规模与主要参与者"]
  }},
  {{
    "heading": "2. 关键技术分析",
    "subheadings": ["2.1 核心技术1", "2.2 核心技术2"]
  }}
]

需要解析的大纲如下：
---
{text}
---

你的JSON输出：
"""
    messages = [{"role": "system",
                 "content": "你是一个专门用于解析文本并输出结构化JSON数据的AI助手。"},
                {"role": "user", "content": prompt}]

    async for chunk in generate_chat_stream(messages):
        full_response += chunk

    return full_response.strip()


async def parse_outline_with_llm(outline: str) -> List[Dict]:
    """
    使用LLM来异步解析大纲，提取包含子标题的嵌套结构。
    """
    logger.info("使用LLM进行深度大纲解析...")
    try:
        json_string = await call_llm_for_parsing(outline)
        # 增加鲁棒性，提取被代码块包裹的JSON
        if "```json" in json_string:
            json_string = json_string.split("```json")[1].split("```")[0]

        start_index = json_string.find('[')
        end_index = json_string.rfind(']')
        if start_index != -1 and end_index != -1:
            json_string = json_string[start_index:end_index + 1]

        headings_structure = json.loads(json_string)
        logger.info(f"LLM成功解析出大纲结构: {headings_structure}")
        return headings_structure
    except Exception as e:
        logger.error(f"使用LLM解析大纲结构时发生异常: {e}", exc_info=True)
        return []


# --- 3. 核心的异步报告生成 Agent (关键改动：分层生成与完整返回) ---
async def generate_report_agent(source_text: str, outline: str, title: str):
    """
    报告生成Agent工作流 (异步生成器)。
    支持嵌套大纲，并为每个子标题生成内容后整体返回。
    """
    try:
        # 异步调用大纲解析
        headings_structure = await parse_outline_with_llm(outline)
        if not headings_structure:
            yield json.dumps({"code": 400, "msg": "Agent无法从大纲中解析出有效的层级结构。", "data": None},
                             ensure_ascii=False)
            return

        full_report_content = ""

        # 关键改动：使用嵌套循环遍历大纲结构
        for main_section in headings_structure:
            main_heading = main_section.get("heading", "")
            subheadings = main_section.get("subheadings", [])

            if not main_heading:
                continue

            # 先将大标题写入最终报告和返回流
            main_heading_md = f"\n\n## {main_heading}\n\n"
            full_report_content += main_heading_md
            yield json.dumps({"code": 200, "msg": "生成中...", "data": {"text": main_heading_md}}, ensure_ascii=False)

            if not subheadings:  # 如果大标题下没有子标题，则为大标题本身生成内容
                subheadings = [main_heading]  # 将大标题视为唯一的子任务

            for subheading in subheadings:
                logger.info(f"Agent开始生成子章节: {subheading}")

                instruction_tmpl = f"""你是一位顶尖的行业分析师，正在撰写关于“{title}”的深度报告。

**报告的完整大纲是：**
{outline}

**到目前为止已完成的内容（供你参考上下文）：**
---
{full_report_content}
---

**你当前的任务是：**
在 **“{main_heading}”** 这个大章节下，为子标题 **“{subheading}”** 撰写详细、专业的内容。

**撰写要求：**
1.  **聚焦当前子标题**：你的输出**只应包含“{subheading}”这一部分的正文**，不要重复标题。
2.  **深度与专业性**：内容需要详实、有深度，至少包含2-3个充实的段落。
3.  **内容丰富**：请结合背景、案例、数据、多维度分析等元素来丰富内容。
4.  **逻辑连贯**：确保内容与上文({full_report_content})自然衔接。

**核心参考资料：**
---
{source_text}
---

现在，请开始为 **“{subheading}”** 撰写正文内容：
"""
                messages = [{"role": "user", "content": instruction_tmpl}]

                # 关键改动：使用缓冲区，将所有文本块拼接完整后再返回
                buffer = ""
                async for chunk in generate_chat_stream(messages):
                    buffer += chunk

                # 拼接子标题和生成的内容
                subheading_content_md = f"### {subheading}\n\n{buffer}\n"

                # 将完整的子标题部分内容存入记忆
                full_report_content += subheading_content_md

                # 将完整的子标题部分内容一次性 yield 出去
                yield json.dumps({"code": 200, "msg": "生成中...", "data": {"text": subheading_content_md}},
                                 ensure_ascii=False)

        yield json.dumps({"code": 200, "msg": "完成", "data": {"text": ""}}, ensure_ascii=False)

    except Exception as e:
        logger.error(f"报告生成Agent发生异常: {str(e)}", exc_info=True)
        yield json.dumps({"code": 500, "msg": f"服务器内部错误: {str(e)}", "data": None}, ensure_ascii=False)

